package pl.edu.agh.cp;

import pl.edu.agh.neural.core.*;
import pl.edu.agh.som.KohonenMap;
import pl.edu.agh.som.Node;

import java.util.List;

public class CounterPropagationNetwork extends BasicNetwork implements ITrainableNetwork
{
    List<? extends ITrainableLayer> trainableLayers;
    List<? extends ITrainableLayerWithTeacherOLD> teachingLayers;

    public CounterPropagationNetwork(InputLayer inputLayer,
                                     List<? extends ITrainableLayer> layers,
                                     List<? extends ITrainableLayerWithTeacherOLD> teachingLayers)
    {
        super(inputLayer, layers);
        trainableLayers = layers;
        this.teachingLayers = teachingLayers;
    }

    @Override
    public void train(double[][] dataVectors, int steps)
    {

        for (int step = 0; step < steps; step++)
        {
            for (double[] dataVector : dataVectors)
            {
                inputLayer.setValues(dataVector);
                for (ITrainableLayer trainableLayer : trainableLayers)
                {
                    trainableLayer.train(step);
                }
            }
        }

    }

    @Override
    public void train(double[][] dataVectors, double[][] teachingVectors, int steps)
    {
        KohonenMap kohonenLayer = (KohonenMap) trainableLayers.get(0);
        GrossbergLayer grossbergLayer = (GrossbergLayer) trainableLayers.get(1);

        for (int step = 0; step < steps; step++)
        {
            int i = 0;
            for (double[] dataVector : dataVectors)
            {
                inputLayer.setValues(dataVector);
                grossbergLayer.setTeachingValues(teachingVectors[i]);

                Node winner = kohonenLayer.trainAndGetWinner(step);
                grossbergLayer.train(step, winner);
                i++;
            }
        }
    }
}
